import cv2
import time
import os
import numpy as np
import json
from ctypes import *
import platform

class SRC_IMG(Structure):
    class DataUnion(Union):
        _fields_ = [("uint16_data", POINTER(c_uint16)),
                   ("uint8_data", POINTER(c_uint8))]

    _fields_ = [("type", c_int), ("len", c_int), ("data", DataUnion)]

class SRC_ALL(Structure):
    _fields_ = [("rgb", SRC_IMG), ("amp", SRC_IMG), ("depth", SRC_IMG)]

def load_config():
    """加载配置文件"""
    with open('config.json', 'r') as f:
        return json.load(f)

def connect_camera(sdk, camera_ip, port, max_retries=10):
    """连接相机"""
        
    # 连接相机
    handle = sdk.api_connect(camera_ip.encode('utf-8'), port)
    if handle <= 0:
        raise ConnectionError(f"无法连接到相机 {camera_ip}:{port}")
    
    print("相机连接成功，等待相机初始化...")
    time.sleep(10)  # 等待相机初始化
    return handle

def connect_camera_with_retry(sdk, camera_ip, port, max_retries=10):
    """尝试连接相机，直到成功或达到最大重试次数"""
    # SDK初始化
    sdk.api_init()
    print("SDK初始化完成")
    
    # 等待一段时间让相机准备就绪
    time.sleep(2)
    print("尝试连接相机...")
    for i in range(max_retries):
        try:
            handle = connect_camera(sdk, camera_ip, port)
            return handle
        except ConnectionError:
            print(f"无法连接到相机，重试中... ({i+1}/{max_retries})")
            time.sleep(1)  # 等待1秒后重试
    
    raise ConnectionError(f"无法连接到相机 {camera_ip}:{port}，已达到最大重试次数")

def get_camera_parameters(sdk, handle, save_path):
    """读取并保存相机参数"""
    # 设置返回类型为字符串（JSON格式）
    sdk.api_get_intrinsic_parameters.restype = c_char_p
    
    # 获取相机参数
    params_json = sdk.api_get_intrinsic_parameters(handle)
    if not params_json:
        raise RuntimeError("无法获取相机参数")
    
    try:
        # 解码并解析JSON数据
        params_str = params_json.decode('utf-8')
        camera_params = json.loads(params_str)
        
        # 保存参数到文件
        params_file = os.path.join(save_path, 'camera_parameters.json')
        with open(params_file, 'w') as f:
            json.dump(camera_params, f, indent=4)
        
        print(f"相机参数已保存到: {params_file}")
        return camera_params
        
    except json.JSONDecodeError as e:
        print(f"JSON解析错误: {e}")
        print(f"原始数据: {params_str}")
        raise
    except Exception as e:
        print(f"处理相机参数时出错: {e}")
        raise

def create_depth_visualization(depth_data, min_depth=0, max_depth=4000):
    """
    创建深度图的可视化版本
    近距离显示为红色，远距离显示为蓝色
    0和65535显示为黑色（无效深度值）
    参数:
        depth_data: 16位深度图
        min_depth: 最小深度值（毫米）
        max_depth: 最大深度值（毫米）
    """
    # 创建掩码，标记有效深度值
    valid_mask = (depth_data != 0) & (depth_data != 65535)
    
    # 创建输出图像（初始为黑色）
    depth_color = np.zeros((depth_data.shape[0], depth_data.shape[1], 3), dtype=np.uint8)
    
    # 只处理有效深度值
    if np.any(valid_mask):
        # 提取有效深度值
        valid_depth = depth_data[valid_mask]
        
        # 裁剪深度值到指定范围
        valid_depth = np.clip(valid_depth, min_depth, max_depth)
        
        # 反转深度值并归一化到0-255
        valid_depth = max_depth - valid_depth
        valid_depth = ((valid_depth - min_depth) / (max_depth - min_depth) * 255).astype(np.uint8)
        
        # 创建彩色映射
        temp_color = cv2.applyColorMap(valid_depth, cv2.COLORMAP_JET)
        
        # 将彩色值填充到有效区域
        depth_color[valid_mask] = temp_color.reshape(-1, 3)[range(np.sum(valid_mask))]
    
    return depth_color

def save_images(img_data, img_type, timestamp, save_path, image_count):
    img_filename = os.path.join(save_path, f'{img_type}_{timestamp}_{image_count}.png')
    cv2.imwrite(img_filename, img_data)          
    return img_filename

def undistort_image(image, camera_params):
    """
    对RGB图像进行畸变矫正
    参数:
        image: 输入的RGB图像
        camera_params: 相机内参
    返回:
        矫正后的图像
    """
    # 提取内参和畸变系数
    camera_matrix = np.array([
        [camera_params["f_x"], 0, camera_params["c_x"]],
        [0, camera_params["f_y"], camera_params["c_y"]],
        [0, 0, 1]
    ])
    dist_coeffs = np.array([
        camera_params["k_1"],
        camera_params["k_2"],
        camera_params["p_1"],
        camera_params["p_2"],
        camera_params["k_3"]
    ])
    
    # 矫正图像
    h, w = image.shape[:2]
    new_camera_matrix, _ = cv2.getOptimalNewCameraMatrix(camera_matrix, dist_coeffs, (w, h), 0, (w, h))
    undistorted_image = cv2.undistort(image, camera_matrix, dist_coeffs, None, new_camera_matrix)
    return undistorted_image

def load_sdk():
    """加载SDK"""
    if platform.system() == 'Linux':
        sdk_path = os.path.join("lib/linux", "libeds_tof_sdk.so")
    elif platform.system() == 'Windows':
        sdk_path = os.path.join("lib\windows", "SDK_WIN.dll")
    else:
        raise OSError("不支持的操作系统")
    
    try:
        sdk = cdll.LoadLibrary(sdk_path)
        print(f"成功加载SDK: {sdk_path}")
    except OSError:
        print(f"错误：无法加载SDK文件 {sdk_path}")
        return None
    
    return sdk

def main():
    # 加载配置
    config = load_config()
    camera_ip = config["camera_ip"]
    port = config["port"]
    save_path = config["save_path"]
    undistort = config.get("undistort", False)  # 获取畸变矫正开关

    # 确保保存目录存在
    os.makedirs(save_path, exist_ok=True)

    # 加载SDK
    sdk = load_sdk()

    # 图像计数器
    handle = -1
    image_count = 0

    # 连接相机
    try:
        handle = connect_camera_with_retry(sdk, camera_ip, port)
    except ConnectionError as e:
        print(f"错误：{e}")
        return
            
    try:

        # 设置返回类型
        sdk.api_get_img.restype = POINTER(SRC_ALL)

        camera_params = None
        
        # 不断尝试获取相机参数，直到成功
        while camera_params is None:
            try:
                camera_params = get_camera_parameters(sdk, handle, save_path)
                print("成功读取相机参数")
            except Exception as e:
                print(f"读取相机参数失败: {e}")
                print("重试获取相机参数...")
                time.sleep(1)  # 等待1秒后重试

        while True:
            # 获取图像
            img = sdk.api_get_img(handle)
            if not bool(img):
                print("警告：无法获取图像")
                time.sleep(0.1)  # 添加小延时
                continue
            
            # 处理RGB图像
            rgb_length = img.contents.rgb.len
            if rgb_length > 0:
                rgb_data = img.contents.rgb.data.uint8_data[:rgb_length]
                rgb_array = np.array(rgb_data, dtype=np.uint8)
                rgb_img = cv2.imdecode(rgb_array, cv2.IMREAD_COLOR)
                    
            # 处理深度图像
            depth_length = img.contents.depth.len // 2  # 16位数据
            depth_data = np.array(list(img.contents.depth.data.uint16_data[:depth_length]), dtype=np.uint16)
            depth_data = depth_data.reshape((480, 640))  # 深度图分辨率为640x480

            # 处理IR图像
            amp_length = img.contents.amp.len // 2  # 16位数据
            amp_data = np.array(list(img.contents.amp.data.uint16_data[:amp_length]), dtype=np.uint16)
            amp_data = amp_data.reshape((480, 640))  # IR图分辨率为640x480

            # 如果启用了畸变矫正，进行矫正
            if undistort and camera_params:
                amp_data = undistort_image(amp_data, camera_params)
                if rgb_length > 0:
                    rgb_img = undistort_image(rgb_img, camera_params["rgb"])

            # 创建深度图的可视化版本
            depth_color = create_depth_visualization(depth_data, min_depth=0, max_depth=4000)
            amp_show = np.sqrt(amp_data * 32)
            amp_show = amp_show.astype(np.uint8)

            # 显示图像
            cv2.imshow("Depth Image", depth_color)
            cv2.imshow("IR Image", amp_show)
            if rgb_length > 0:
                cv2.imshow("RGB Image", rgb_img)

            # 处理键盘事件
            key = cv2.waitKey(1)
            if key == ord('s'):  # 按's'保存图像
                timestamp = time.strftime('%Y%m%d_%H%M%S', time.localtime())
                image_count += 1
                if rgb_length > 0:
                    depth_filename = save_images(depth_data, 'depth', timestamp, save_path, image_count)
                    ir_filename = save_images(amp_data, 'ir', timestamp, save_path, image_count)
                    rgb_filename = save_images(rgb_img, 'rgb', timestamp, save_path, image_count)
                    print(f"图像已保存:")
                    print(f"IR: {ir_filename}")
                    print(f"Depth: {depth_filename}")
                    print(f"RGB: {rgb_filename}")
                else:
                    depth_filename = save_images(depth_data, 'depth', timestamp, save_path, image_count)
                    ir_filename = save_images(amp_data, 'ir', timestamp, save_path, image_count)
                    print(f"图像已保存:")
                    print(f"IR: {ir_filename}")
                    print(f"Depth: {depth_filename}")
            elif key == ord('q'):  # 按'q'退出
                print("正在退出程序...")
                break

    finally:
        # 清理资源
        if handle > 0:
            sdk.api_disconnect(handle)
        cv2.destroyAllWindows()
        print("相机已断开连接")

if __name__ == '__main__':
    main()